from typing import Optional, Dict

import chromadb
from chromadb import Where, WhereDocument, QueryResult, Settings
from chromadb.api.models.Collection import Collection
from chromadb.errors import NotFoundError

from app.config.logging import logger
from app.decorator.timeit import timeit
from app.llm.ollama import bge_embedding_fun

# 创建chroma持久化客户端
# 启用anonymized_telemetry 匿名数据收集功能，用来检测chroma的操作
client = chromadb.PersistentClient(path="D:\\temp\\chroma", settings=Settings(anonymized_telemetry=True))


class ChromaClient:

    # 可以默认指定1个embedding_func
    def __init__(self, collection_name: str, embedding_func=bge_embedding_fun) -> None:
        self.collection_name = collection_name
        self.embedding_func = embedding_func
        self.client = client
        self.mode = 'local'
        self.collection = self._get_collection(collection_name)

    def _get_collection(self, collection_name: str) -> Collection:
        """
        获取Collection集合
        :param collection_name: 集合名称
        :return:
        """
        try:
            collection = self.client.get_collection(collection_name)
        except NotFoundError:
            logger.info(f"Collection {collection_name} does not exist. will be create collection.")
            collection = self.client.create_collection(name=collection_name, embedding_function=self.embedding_func,
                                                       metadata={"title": "知识库"}, get_or_create=True)
        return collection

    @timeit
    def upsert_documents(self, documents: list[str], ids: list[str] = None,
                         metadata: Optional[list[Dict]] = None, ) -> None:
        """
        添加document到指定知识库
        """
        # embedding 是 内容向量化之后的
        self.collection.upsert(ids=ids, metadatas=metadata, documents=documents)

    @timeit
    def search_documents(self, query: str, ids: list[str] = None, where: Optional[Where] = None,
                         where_document: Optional[WhereDocument] = None,
                         top_k: int = 3) -> QueryResult:
        """
        搜索
        :param where:
        :param where_document:
        :param top_k:
        :param query:
        :param ids:
        :param metadata:
        :return:
        """
        return self.collection.query(query_texts=query, ids=ids, n_results=top_k, where=where,
                                     where_document=where_document)
